from langchain.vectorstores import Chroma
from langchain.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OllamaEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import Ollama
from langchain.prompts import PromptTemplate
import sys



# 指定要加载的网页URL
url = "https://cloud.tencent.com/developer/article/2508298"

# 创建WebBaseLoader实例
loader = WebBaseLoader(url)
# 加载文档
data = loader.load()

# 假设`data`是我们从前面步骤中加载的文档数据
# 使用RecursiveCharacterTextSplitter对文档进行分割
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)

# 执行分割操作
all_splits = text_splitter.split_documents(data)

# 打印分割后的文档片段数量，以及前一个片段的内容，以验证分割是否成功
print(f"总共分割出{len(all_splits)}个片段。")
# print(f"第一个片段的内容为：{all_splits[0]}")
# for item in all_splits:
#     print("\n")
#     print(f"-------------------------{item}")

# 初始化Ollama嵌入模型
oembed = OllamaEmbeddings(base_url="http://localhost:11434", model="nomic-embed-text")

# 为每个文档片段生成向量
vectors = [
    oembed.embed_documents([doc.page_content])[0]  # 提取文档内容并生成向量
    for doc in all_splits
]

# 打印出第一个文档片段的向量，以验证向量生成是否成功
print(f"片段向量个数={len(vectors)}")
# print(vectors[0])

# 使用文档片段构建ChromaDB向量数据库
vectorstore = Chroma.from_documents(
    documents=all_splits,  # 原始文档
    embedding=oembed       # 嵌入模型
)

# 现在，您可以使用vectorstore来进行相似性搜索
# 假设用户的查询是
query = "什么是智能体Agent"

# 将查询转换为向量
query_vector = oembed.embed_query(query)

# 在向量数据库中搜索最相似的文档片段
similar_documents = vectorstore.similarity_search(query_vector)

# 初始化Ollama语言模型 (修正后)
ollama = Ollama(
    base_url="http://localhost:11434",
    model="qwen2.5:7b"
)

template = """使用以下上下文来回答最后的问题。如果你不知道答案，只需说你不知道，不要试图编造答案。
        上下文: {context}
        问题: {question}
        回答: """
QA_PROMPT = PromptTemplate(
        template=template,
        input_variables=["context", "question"]
    )

# 创建QA链
qachain = RetrievalQA.from_chain_type(
    llm=ollama,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    return_source_documents=True,
    chain_type_kwargs={"prompt": QA_PROMPT}
)

# 使用RetrievalQA回答问题
# answer = qachain.invoke({"query": "什么是智能体Agent？"})

# 创建带进度提示的流式处理
answer_stream = qachain.stream({"query": "什么是智能体Agent？"})

print("思考中", end="")
for chunk in answer_stream:
    if "result" in chunk:
        print("\r回答：", end="")  # 清除进度提示
        for char in chunk["result"]:
            print(char, end='', flush=True)
            sys.stdout.flush()  # 确保立即输出
    # 处理文档来源
    if "source_documents" in chunk:
        print("\n\n=== 参考来源 ===")
        for i, doc in enumerate(chunk["source_documents"], 1):
            print(f"[文档{i}] {doc.metadata['source']}")
            print(f"内容摘要: {doc.page_content[:120]}...\n")
